package org.skipper.spring.cloud.dubbo.helper;

import lombok.extern.log4j.Log4j2;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import org.apache.logging.log4j.ThreadContext;
import org.skipper.common.constant.HeaderKeys;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.UUID;

import static org.apache.dubbo.common.constants.CommonConstants.CONSUMER;

/**
 * Created with IntelliJ IDEA.
 * Description:
 *
 * @author shoujian
 * Date: 2022-02-14
 * Time: 11:42 AM
 */
@Activate(group = {CONSUMER})
@Log4j2
public class SpringCloudConsumerDubboContextFilter implements Filter {

    private static final String TRACE_ID = "traceId";

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {


        RpcContext context = RpcContext.getContext();

        context.setInvoker(invoker)
                .setInvocation(invocation)
                .setLocalAddress(invoker.getUrl().getHost(), invoker.getUrl().getPort())
        ;

        RequestAttributes requestAttributes = getRequestAttributesSafely();

        if (requestAttributes != null) {
            ServletRequestAttributes sra = (ServletRequestAttributes) requestAttributes;
            HttpServletRequest request = sra.getRequest();
            for (HeaderKeys headerKeys : HeaderKeys.values()) {
                context.setAttachment(headerKeys.getHeaderKey(), request.getHeader(headerKeys.getHeaderKey()));
            }
        }

        String traceId = UUID.randomUUID().toString();
        // 存入 Log4j2 ThreadContext
        ThreadContext.put(TRACE_ID, traceId);
        context.setAttachment(TRACE_ID, traceId);

        if (invocation instanceof RpcInvocation) {
            ((RpcInvocation) invocation).setInvoker(invoker);
        }

        try {
            context.clearAfterEachInvoke(false);
            return invoker.invoke(invocation);
        } finally {
            context.clearAfterEachInvoke(true);
            // IMPORTANT! For async scenario, we must remove context from current thread, so we always create a new RpcContext for the next invoke for the same thread.
            RpcContext.removeContext(true);
            RpcContext.removeServerContext();
        }

    }


    public RequestAttributes getRequestAttributesSafely(){
        RequestAttributes requestAttributes;
        try{
            requestAttributes = RequestContextHolder.currentRequestAttributes();
        }catch (IllegalStateException e){
            requestAttributes = null;
        }
        return requestAttributes;
    }
}
